package com.azure.code.graph.prim;

import java.util.List;
import java.util.PriorityQueue;

/**
 * 最小生成树算法：
 * Kruska 算法用到了 Union-Find并查集算法，来保证挑选出来的这些边组成的一定是一颗【树】，而不会包含环或者形成
 *        一片森林
 * Prim
 *     首先，Prim算法也是使用贪心思想来让生成树的权重仅可能小   切分定理
 *     其次，Prim算法使用BFS算法思想和visited布尔数组避免成环，来保证选出的边最终形成的一定是一颗树
 * Prim算法不需要事先对所有的边排序，而是利用优先级队列动态事先排序效果
 *
 * 切分定理：
 *     对任意一种切分，其中权重最小的那条 横切边 一定是构成最小生成树的一条边
 *
 * Prim算法的逻辑： 每次切分都能找到最小生成树的一条边，然后可以进行新一轮切分，直到找到最小生成树的所有边为止
 */
public class Prim {
    // 核心数据结构，存储【横切边】的优先级队列
    private PriorityQueue<int[]> pq;
    // 类似 visited数组的作用，记录哪些节点已经成为最小生成树的一部分
    private boolean[] inMst;
    // 记录最小生成树的权重和
    private  int weightSum = 0;
    // graph 是用邻接表 表示的一幅图
    // graph[s] 记录节点 s 所有的相邻的边
    // 三元组 int[]{from,to,weight}表示一条边
    private List<int[]>[] graph;

    public Prim(List<int[]>[] graph) {
        this.graph = graph;
        // 按照边的权重 从小到大 排序
        this.pq = new PriorityQueue<>((a,b)-> a[2]-b[2]);
        // 图中有n个节点
        int n = graph.length;
        this.inMst = new boolean[n];

        // 随便从一个点开始切分都可以，我们不妨从节点0开始
        inMst[0] = true;
        cut(0);
        // 不断切分，向最小生成树种添加边
        while (!pq.isEmpty()){
            int[] edge = pq.poll();
            int to = edge[1];
            int weight = edge[2];
            if(inMst[to]){
                // 节点to已经在最小生成树种，跳过
                // 否则这条边会产生环
                continue;
            }
            // 将边edge加入最小生成树中
            weightSum+=weight;
            inMst[to] = true;
            // 节点to加入后，进行新一轮的切分，会产生更多横切边
            cut(to);
        }
    }
    // 将 s 的横切边加入到优先队列中
    private void cut(int s) {
        // 遍历 s 的邻边
        for (int[] edge: graph[s]){
            int to  = edge[1];
            if(inMst[to]){
                // 相邻节点to已经在最小生成树中，跳过
                // 否则这条边会产生环
                continue;
            }
            // 加入横切边队列
            pq.offer(edge);
        }
    }

    // 最小生成树的权重和
    public int getWeightSum(){
        return weightSum;
    }

    // 判断最小生成树是否包含图汇总的所有节点
    public boolean allConnected(){
        for(int i = 0; i < inMst.length; i++){
            if(!inMst[i]){
                return false;
            }
        }
        return true;
    }
}
